import numpy as np
import matplotlib.pyplot as plt

def fftfreq2d(m,n=None,d=1,zero='center',axis=0):
    x=np.fft.fftfreq(m,d)
    if n==None: n=m
    y=np.fft.fftfreq(n,d)
    X,Y=np.meshgrid(x,y)
    if axis==0: result=X
    else: result=Y
    if zero=='center':
        result=np.fft.fftshift(result,1-axis)
    return result

def shiftimg(im_fft,shifts,zero='front'):
    """fft of image -> shifted fft of image by x=shifts[0], y=shifts[1]"""
    im_fft_shift=im_fft[:]
    for n in range(len(shifts)):
        im_fft_shift=im_fft_shift*np.exp(1j*fftfreq2d(np.shape(im_fft)[1],np.shape(im_fft)[0],axis=n,zero=zero)*2*np.pi*(-shifts[n]))
    return im_fft_shift

def zeropadding(im,s):
    """Pad ndarray im with zeros arround the original ndarray.
    Number of zeros (on each side) padded for different dimensions are declared in s.
    If the length of s is shorter than the dimensions of im, s will be padded with
    its last value in the end."""
    dim=len(np.shape(im))
    if type(s)==int:
        s=[s]
    if len(s)<dim:
        s0=s[-1]
        for n in range(len(s),dim):
            s=np.append(s,s0)
    im_padded=im[:]
    for n in range(dim):
        shape=list(np.shape(im_padded))
        shape[n]=s[n]
        append=np.zeros(shape)
        im_padded=np.append(append,im_padded,axis=n)
        im_padded=np.append(im_padded,append,axis=n)
    return im_padded

#test pattern    
a=np.arange(5)
a=np.append(a,a[::-1])
b=np.copy(a)
x,y=np.meshgrid(a,b)
im=x*y


#image
im=plt.imread('t.jpg')
im=np.mean(im,2)
im_fft=np.fft.fft2(im)



"""
#upsampling
im_fft=np.fft.fftshift(im_fft,axes=[0,1])
im_fft=zeropadding(im_fft,10)
im_fft=np.fft.ifftshift(im_fft,axes=[0,1])
"""


im_fft_shift=shiftimg(im_fft,[53.5,0])

im_1=np.fft.ifft2(im_fft_shift)
plt.imshow(np.real(im_1))
plt.show()